import sympy as sym
import numpy as np
import math

import sympy as sym
import numpy as np
from sympy import sin,cos,Function
import numbers
sym.init_printing()
#document:https://github.com/sympy/sympy


def iterable(arg):
    T = True
    try:
        for i in arg:
            pass
    except TypeError as err:
        T = False
    return T

def nthroot(arg_a, arg_b, flag = True):
    res = []
    if iterable(arg_a) and iterable(arg_b):
        if x.shape == n.shape:
            rows, cols = arg_a.shape 
            res = np.empty(arg_a.shape, dtype=complex)             
            for i in range(rows):               
                for j in range(cols):           
                    if flag:
                        res[i, j] = sym.simplify(sym.root(arg_a[i, j], arg_b[i, j]))
                    else:
                        res[i, j] = sym.root(arg_a[i, j], arg_b[i, j])
            return res
        else:
            raise TypeError("Input type not supported.") 
    elif (iterable(arg_a) and not iterable(arg_b)) or (not iterable(arg_a) and iterable(arg_b)):
        raise TypeError("Input type not supported.") 
    else:
        if flag:
            return sym.simplify(sym.root(arg_a, arg_b))
        else:
            return sym.root(arg_a, arg_b)






print(nthroot(27, 3))
print(nthroot(sym.E**3,sym.E))
print(nthroot(8, 2))

x = np.array([[-27,-8,-4],
    [27,64,-12]])

n = np.array([[3,3,4],
    [3,2,-2]])

print(nthroot(x, n))

print(nthroot(-1/3, -3))
print(nthroot(-27, -3))
print(sym.simplify(sym.root(-27,-3)))
